# 看Norm的覆盖范围————思路：找平衡数值
import torch
import torch.nn as nn

from einops import rearrange, reduce

torch.random.manual_seed(0)

# data generating
x = torch.tensor([[[[1, 2, 3], [1, 2, 3], [1, 2, 3]],
                   [[10, 20, 30], [10, 20, 30], [10, 20, 30]],
                   [[100, 200, 300], [100, 200, 300], [100, 200, 300]]],
                  [[[4, 5, 6], [4, 5, 6], [4, 5, 6]],
                   [[40, 50, 60], [40, 50, 60], [40, 50, 60]],
                   [[400, 500, 600], [400, 500, 600], [400, 500, 600]]],
                  [[[7, 8, 9], [7, 8, 9], [7, 8, 9]],
                   [[70, 80, 90], [70, 80, 90], [70, 80, 90]],
                   [[700, 800, 900], [700, 800, 900], [700, 800, 900]]]],
                 dtype=torch.float32)

x = torch.cat([x, x + 0.1], dim=2)
x = rearrange(x, 'b c (n h) w -> b (c n) h w', n=2)

# data shape measurement
b, c, h, w = x.shape
print('x:', x.shape)
print(x)

# data statistics
for i in range(b):
    print(x[i].mean())

for i in range(c):
    print(x[:, i, :, :].mean())

for i in range(b):
    for j in range(c):
        print(x[i, j, :, :].mean())

# grouping channels into how many pieces
group_num = 3

# display
print(50 * '-', 'BN', 50 * '-')
bn = nn.BatchNorm2d(c, eps=1e-10, affine=False, track_running_stats=False)
y_bn = bn(x)
print(y_bn)

print(50 * '-', 'LN', 50 * '-')
ln = nn.LayerNorm([c, h, w], eps=1e-12, elementwise_affine=False)
y_ln = ln(x)
print(y_ln)

print(50 * '-', 'IN', 50 * '-')
In = nn.InstanceNorm2d(c, eps=1e-12, affine=False, track_running_stats=False)
y_In = In(x)
print(y_In)

print(50 * '-', 'GN', 50 * '-')
gn = nn.GroupNorm(group_num, c, eps=1e-12, affine=False)
y_gn = gn(x)
print(y_gn)
